package org.eclipse.scanning.test.scan.nexus;

import static org.eclipse.scanning.sequencer.analysis.ClusterProcessingRunnableDevice.DEFAULT_ENTRY_PATH;
import static org.eclipse.scanning.sequencer.analysis.ClusterProcessingRunnableDevice.NEXUS_FILE_EXTENSION;
import static org.eclipse.scanning.sequencer.nexus.SolsticeConstants.GROUP_NAME_SOLSTICE_SCAN;
import static org.eclipse.scanning.test.scan.nexus.NexusAssert.assertAxes;
import static org.eclipse.scanning.test.scan.nexus.NexusAssert.assertIndices;
import static org.eclipse.scanning.test.scan.nexus.NexusAssert.assertSignal;
import static org.eclipse.scanning.test.scan.nexus.NexusAssert.assertTarget;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.CoreMatchers.instanceOf;
import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.CoreMatchers.nullValue;
import static org.hamcrest.CoreMatchers.startsWith;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertSame;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.assertTrue;

import java.util.Arrays;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;

import org.eclipse.dawnsci.analysis.api.tree.DataNode;
import org.eclipse.dawnsci.analysis.api.tree.TreeFile;
import org.eclipse.dawnsci.nexus.INexusFileFactory;
import org.eclipse.dawnsci.nexus.NXdata;
import org.eclipse.dawnsci.nexus.NXdetector;
import org.eclipse.dawnsci.nexus.NXentry;
import org.eclipse.dawnsci.nexus.NXinstrument;
import org.eclipse.dawnsci.nexus.NXpositioner;
import org.eclipse.dawnsci.nexus.NXroot;
import org.eclipse.dawnsci.nexus.NexusFile;
import org.eclipse.dawnsci.nexus.NexusUtils;
import org.eclipse.january.dataset.IDataset;
import org.eclipse.january.dataset.PositionIterator;
import org.eclipse.scanning.api.device.AbstractRunnableDevice;
import org.eclipse.scanning.api.device.IRunnableDevice;
import org.eclipse.scanning.api.event.core.IConsumer;
import org.eclipse.scanning.api.event.scan.DeviceState;
import org.eclipse.scanning.api.event.status.StatusBean;
import org.eclipse.scanning.api.points.IPosition;
import org.eclipse.scanning.api.scan.models.ScanModel;
import org.eclipse.scanning.test.scan.mock.DummyOperationBean;

public class ScanClusterProcessingChecker {

	private INexusFileFactory factory;
	private IConsumer<StatusBean> consumer;
	private List<String> scannableNames;
	private String filePath;
	private String detectorName;
	private String processingName = "sum";

	public ScanClusterProcessingChecker(INexusFileFactory factory, IConsumer<StatusBean> consumer) {
		this.factory = factory;
		this.consumer= consumer;
	}

	public void setDevice(IRunnableDevice<ScanModel> scanner) throws Exception {

		final ScanModel scanModel = ((AbstractRunnableDevice<ScanModel>) scanner).getModel();

		final IPosition pos = scanModel.getPositionIterable().iterator().next();
		this.scannableNames = pos.getNames();
		this.filePath = ((AbstractRunnableDevice<ScanModel>) scanner).getModel().getFilePath();
		this.detectorName = scanModel.getDetectors().get(0).getName();

		assertEquals(DeviceState.ARMED, scanner.getDeviceState());
	}

	public void checkNexusFile(int... sizes) throws Exception {


		NXroot rootNode = getNexusRoot();
		NXentry entry = rootNode.getEntry();
		NXinstrument instrument = entry.getInstrument();

		LinkedHashMap<String, List<String>> signalFieldAxes = new LinkedHashMap<>();
		// axis for additional dimensions of a datafield, e.g. image
		signalFieldAxes.put(NXdetector.NX_DATA, Arrays.asList("real", "imaginary"));
		signalFieldAxes.put("spectrum", Arrays.asList("spectrum_axis"));
		signalFieldAxes.put("value", Collections.emptyList());

		NXdetector detector = instrument.getDetector(detectorName);
		// map of detector data field to name of nxData group where that field is the @signal field
		Map<String, String> expectedDataGroupNames =
				signalFieldAxes.keySet().stream().collect(Collectors.toMap(Function.identity(),
				x -> detectorName + (x.equals(NXdetector.NX_DATA) ? "" : "_" + x)));

		// validate the main NXdata generated by the NexusDataBuilder
		Map<String, NXdata> nxDataGroups = entry.getChildren(NXdata.class);
		assertEquals(signalFieldAxes.size(), nxDataGroups.size());
		assertTrue(nxDataGroups+" should have these in: "+expectedDataGroupNames, nxDataGroups.keySet().containsAll(
				expectedDataGroupNames.values()));
		for (String nxDataGroupName : nxDataGroups.keySet()) {
			NXdata nxData = entry.getData(nxDataGroupName);

			String sourceFieldName = nxDataGroupName.equals(detectorName) ? NXdetector.NX_DATA :
				nxDataGroupName.substring(nxDataGroupName.indexOf('_') + 1);
			assertSignal(nxData, sourceFieldName);
			// check the nxData's signal field is a link to the appropriate source data node of the detector
			DataNode dataNode = detector.getDataNode(sourceFieldName);
			IDataset dataset = dataNode.getDataset().getSlice();
			assertSame(dataNode, nxData.getDataNode(sourceFieldName));
			assertTarget(nxData, sourceFieldName, rootNode, "/entry/instrument/" + detectorName
					+ "/" + sourceFieldName);

			// check that the other primary data fields of the detector haven't been added to this NXdata
			for (String primaryDataFieldName : signalFieldAxes.keySet()) {
				if (!primaryDataFieldName.equals(sourceFieldName)) {
					assertNull(nxData.getDataNode(primaryDataFieldName));
				}
			}

			int[] shape = dataset.getShape();
			for (int i = 0; i < sizes.length; i++)
				assertEquals(sizes[i], shape[i]);

			// Make sure none of the numbers are NaNs. The detector
			// is expected to fill this scan with non-nulls.
			final PositionIterator it = new PositionIterator(shape);
			while (it.hasNext()) {
				int[] next = it.getPos();
				assertFalse(Double.isNaN(dataset.getDouble(next)));
			}

			// Check axes

			// Append _value_demand to each name in list, then add detector axis fields to result
			List<String> expectedAxesNames = Stream.concat(
					scannableNames.stream().map(x -> x + "_value_set"),
					signalFieldAxes.get(sourceFieldName).stream()).collect(Collectors.toList());
			assertAxes(nxData, expectedAxesNames.toArray(new String[expectedAxesNames.size()]));

			int[] defaultDimensionMappings = IntStream.range(0, sizes.length).toArray();
			int i = -1;
			for (String  scannableName : scannableNames) {

			    i++;
				NXpositioner positioner = instrument.getPositioner(scannableName);
				assertNotNull(positioner);

				dataNode = positioner.getDataNode("value_set");
				dataset = dataNode.getDataset().getSlice();
				shape = dataset.getShape();
				assertEquals(1, shape.length);
				assertEquals(sizes[i], shape[0]);

				String nxDataFieldName = scannableName + "_value_set";
				assertSame(dataNode, nxData.getDataNode(nxDataFieldName));
				assertIndices(nxData, nxDataFieldName, i);
				assertTarget(nxData, nxDataFieldName, rootNode,
						"/entry/instrument/" + scannableName + "/value_set");

				// Actual values should be scanD
				dataNode = positioner.getDataNode(NXpositioner.NX_VALUE);
				dataset = dataNode.getDataset().getSlice();
				shape = dataset.getShape();
				assertArrayEquals(sizes, shape);

				nxDataFieldName = scannableName + "_" + NXpositioner.NX_VALUE;
				assertSame(dataNode, nxData.getDataNode(nxDataFieldName));
				assertIndices(nxData, nxDataFieldName, defaultDimensionMappings);
				assertTarget(nxData, nxDataFieldName, rootNode,
						"/entry/instrument/" + scannableName + "/"
								+ NXpositioner.NX_VALUE);
			}
		}
	}


	public void checkSubmittedBean(boolean allowEmpty) throws Exception {

		List<StatusBean> statusSet = consumer.getStatusSet();
		if (allowEmpty && statusSet.isEmpty()) return;

		assertThat(statusSet.get(0), is(instanceOf(DummyOperationBean.class)));
		DummyOperationBean operationBean = (DummyOperationBean) statusSet.get(0);

		assertThat(operationBean.getDataKey(), is(equalTo(DEFAULT_ENTRY_PATH + GROUP_NAME_SOLSTICE_SCAN)));
		assertTrue(operationBean.getFilePath().contains("test_nexus"));
		assertTrue(operationBean.getFilePath().endsWith(NEXUS_FILE_EXTENSION));
		assertTrue(operationBean.getOutputFilePath().contains("processed"));
		assertTrue(operationBean.getOutputFilePath(),
				   operationBean.getOutputFilePath().endsWith("-"+processingName+NEXUS_FILE_EXTENSION));
		assertThat(operationBean.getDatasetPath(), is(equalTo(DEFAULT_ENTRY_PATH + detectorName)));
		assertThat(operationBean.getSlicing(), is(nullValue()));
		assertThat(operationBean.getProcessingPath(), is(equalTo("/tmp/sum.nxs")));
		assertThat(operationBean.isDeleteProcessingFile(), is(false));
		assertThat(operationBean.getAxesNames(), is(nullValue()));
		assertThat(operationBean.getXmx(), is(equalTo("1024m")));
		assertThat(operationBean.getDataDimensions(), is(nullValue()));
		assertThat(operationBean.getScanRank(), is(2));
		assertThat(operationBean.isReadable(), is(true));

		assertThat(System.getProperty("java.io.tmpdir"), startsWith(operationBean.getRunDirectory()));
		assertThat(operationBean.getNumberOfCores(), is(1));
	}

	public NXroot getNexusRoot() throws Exception {

		NexusFile nf = factory.newNexusFile(filePath);
		nf.openToRead();

		TreeFile nexusTree = NexusUtils.loadNexusTree(nf);
		return (NXroot) nexusTree.getGroupNode();
	}

	public List<String> getScannableNames() {
		return scannableNames;
	}

	public void setScannableNames(List<String> scannableNames) {
		this.scannableNames = scannableNames;
	}

	public String getFilePath() {
		return filePath;
	}

	public void setFilePath(String filePath) {
		this.filePath = filePath;
	}

	public String getDetectorName() {
		return detectorName;
	}

	public void setDetectorName(String detectorName) {
		this.detectorName = detectorName;
	}

	public String getProcessingName() {
		return processingName;
	}

	public void setProcessingName(String processingName) {
		this.processingName = processingName;
	}

}
